# models/time_llm_block.py
import paddle
import paddle.nn as nn
from models.autocorrelation import AutoCorrelationMechanism
from models.series_decomp import SeriesDecomposition

class TimeLLMBlock(nn.Layer):
    """融合自相关和序列分解的核心块"""

    def __init__(self, d_model, top_k=3, kernel_size=25):
        super().__init__()
        self.autocorr = AutoCorrelationMechanism(d_model, top_k)
        self.decomp1 = SeriesDecomposition(kernel_size)
        self.decomp2 = SeriesDecomposition(kernel_size)

        # 轻量级FFN
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model)
        )

    def forward(self, x):
        # 自相关处理
        attn_out = self.autocorr(x)
        x = x + attn_out

        # 序列分解
        trend1, seasonal1 = self.decomp1(x)

        # FFN处理季节部分
        ffn_out = self.ffn(seasonal1)
        seasonal2 = seasonal1 + ffn_out

        # 二次分解
        trend2, seasonal_out = self.decomp2(seasonal2)

        # 趋势累加
        trend_out = trend1 + trend2
        return trend_out, seasonal_out
